package cn.veliger.springai.controller;

import cn.veliger.springai.advisor.CusCallAroundAdvisor;
import cn.veliger.springai.entity.ActorFilms;
import lombok.RequiredArgsConstructor;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.chat.prompt.PromptTemplate;
import org.springframework.ai.content.Media;
import org.springframework.ai.converter.BeanOutputConverter;
import org.springframework.core.ParameterizedTypeReference;
import org.springframework.core.io.ClassPathResource;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.util.MimeTypeUtils;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import reactor.core.publisher.Flux;

import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.List;
import java.util.Map;

/**
 * @author veliger
 * create on 2025-06-28
 **/
@RestController
@RequestMapping
@RequiredArgsConstructor
public class ChatController {

    private final ChatClient chatClient;


    private final ChatModel chatModel;

    private final ChatClient memoryChatClient;

    /**
     * 使用chatClient进行聊天
     *
     * @param msg
     * @return
     */
    @GetMapping("/chatClient")
    public String chatClient(String msg) {
        return this.chatClient.prompt()
                .advisors(new CusCallAroundAdvisor())
                .user(msg)
                .call()
                .content();
    }


    /**
     * 使用chatModel进行聊天
     *
     * @param msg
     * @return
     */
    @GetMapping("/chatModel")
    public String chatModel(String msg) {
        ChatClient.Builder builder = ChatClient.builder(chatModel);
        ChatClient chatClient = builder.build();
        return chatClient.prompt().user(msg).call().content();
    }

    /**
     * 使用flux进行聊天
     *
     * @return
     */
    @GetMapping(value = "/streamChat", produces = "text/html;charset=UTF-8")
    public Flux<String> streamChat(String msg) {
        return chatClient.prompt().user(msg).stream().content();
    }

    /**
     * 使用sse
     *
     * @param msg
     * @return
     */
    @GetMapping(value = "/sseChat")
    public SseEmitter sseChat(String msg) {
        SseEmitter sseEmitter = new SseEmitter() {
            @Override
            protected void extendResponse(ServerHttpResponse outputMessage) {
                HttpHeaders headers = outputMessage.getHeaders();
                headers.setContentType(new MediaType("text", "event-stream", StandardCharsets.UTF_8));
            }
        };

        Flux<String> stream = chatClient.prompt().user(msg).stream().content();

        stream.subscribe(token -> {
            try {
                sseEmitter.send(token);
            } catch (IOException e) {
                sseEmitter.completeWithError(e);
            }
        }, sseEmitter::completeWithError, sseEmitter::complete);
        return sseEmitter;
    }

    /**
     * 使用系统提示词
     *
     * @return
     */
    @GetMapping("/systemChat")
    public String systemChat(String msg) {
        return this.chatClient.prompt().system("请用我提供的主题，写一首七言绝句").user(msg).call().content();
    }

    /**
     * 结构化输出
     *
     * @return
     */
    @GetMapping("/entity")
    public ActorFilms entity() {
        return this.chatClient.prompt().user("给我周星驰的三部电影信息").call().entity(ActorFilms.class);
    }

    /**
     * 使用底层的结构化
     *
     * @return
     */
    @GetMapping("/output")
    public List<ActorFilms> output() {
        BeanOutputConverter<List<ActorFilms>> outputConverter = new BeanOutputConverter<>(new ParameterizedTypeReference<List<ActorFilms>>() {
        });

        String format = outputConverter.getFormat();
        String template = """
                给我周星驰和成龙的三部电影信息,返回集合数据
                {format}
                """;
        Prompt prompt = PromptTemplate.builder().template(template).variables(Map.of("format", format)).build().create();
        String content = chatClient.prompt(prompt).call().content();

        return outputConverter.convert(content);
    }

    /**
     * 使用Advisor
     *
     * @param msg
     * @return
     */
    @GetMapping("/advisor")
    public String advisor(String msg) {
        ChatClient chatClient = ChatClient.builder(chatModel).defaultAdvisors(new CusCallAroundAdvisor()).build();
        return chatClient.prompt().user(msg).call().content();
    }

    /**
     * 会话记忆
     *
     * @param chatId
     * @param msg
     * @return
     */
    @GetMapping("/chatMemoryChat")
    public String chatMemoryChat(@RequestParam("chatId") String chatId, @RequestParam("msg") String msg) {
        return memoryChatClient
                .prompt()
                .advisors(a -> a.param(ChatMemory.CONVERSATION_ID, chatId))
                .user(msg)
                .call()
                .content();
    }

    /**
     * 多模态聊天
     *
     * @return
     */
    @GetMapping("/multimodalChat")
    public String multimodalChat() {
        ClassPathResource imageData = new ClassPathResource("images/1.png");
        UserMessage userMessage = UserMessage.builder()
                .text("图片里有什么?")
                .media(Media.builder().data(imageData).mimeType(MimeTypeUtils.IMAGE_PNG).build())
                .build();
        return chatModel.call(userMessage);
    }

}
